import cv2
import numpy as np
import requests
import json
import re
# import gradio as gr
import os
from skimage.morphology import skeletonize
import matplotlib.pyplot as plt
from io import BytesIO
import base64
from PIL import Image
import math

"""框架图识别返回墙体和门窗等"""
class WireFrames:
    def __init__(self):
        pass
    """1.url识别转换"""
    def extract_file_name(self,url):
        match = re.search(r'standard_(.*\.png)', url)
        if match:
            return match.group(1)
        else:
            raise ValueError("File name extraction failed")
    
    def generate_urls(self,file_name):
        if file_name:
            wireframes_url = f"https://vrlab-public.ljcdn.com/release/vrcustomer/outline_0_{file_name}"
            return wireframes_url
        else:
            raise ValueError("File name is not extracted")
    
    def download_image(self,url):
        try:
            response = requests.get(url)
            response.raise_for_status()  # Raise HTTPError for bad response status
            image_array = np.asarray(bytearray(response.content), dtype=np.uint8)
            image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
            if image is None:
                raise ValueError("Failed to decode image")
            return image
        except requests.exceptions.RequestException as e:
            print(f"Error downloading image from {url}: {e}")
            return None
        except Exception as e:
            print(f"Error processing image: {e}")
            return None
       
    """4.线框图wireframe识别"""
    def wireframes_image(self,image,drawWidth=1):
        result = {}
        """ wireframe框线图识别 """
        # image = self.download_image(self.wireframes_url)
        # Convert image to grayscale
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 二值化处理
        _, binary = cv2.threshold(gray_image, 240, 255, cv2.THRESH_BINARY)

        # 检测轮廓
        contours, _ = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        # 创建一个空白图像用于绘制检测结果
        output = cv2.cvtColor(gray_image, cv2.COLOR_GRAY2BGR)

        # 绘制轮廓
        cv2.drawContours(output, contours, -1, (0, 255, 0), drawWidth)

        # 将结果图像保存到内存中
        _, buffer = cv2.imencode('.jpg', output)
        detected_contours = cv2.imdecode(np.frombuffer(buffer, np.uint8), cv2.IMREAD_COLOR)

        # 转换为灰度图像
        gray_wall_only2 = cv2.cvtColor(detected_contours, cv2.COLOR_BGR2GRAY)

        # 二值化处理
        _, binary_image2 = cv2.threshold(gray_wall_only2, 127, 255, cv2.THRESH_BINARY)

        # 骨架化
        skeleton2 = skeletonize(binary_image2 // 255)
        skeleton2 = (skeleton2 * 255).astype(np.uint8)
        skeleton_image2 = np.zeros_like(detected_contours)
        skeleton_image2[skeleton2 == 255] = [255, 255, 255]

        # 检测骨架图像的轮廓
        skeleton_contours2, _ = cv2.findContours(skeleton2, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        linedoor = []
        for i, contour in enumerate(skeleton_contours2):
            for j in range(len(contour)):
                x1, y1 = contour[j][0]
                x2, y2 = contour[(j + 1) % len(contour)][0]
                linedoor.append((1, x1, y1, x2, y2))

        # Detect white and black regions
        lower_white = np.array([200, 200, 200])
        upper_white = np.array([255, 255, 255])
        white_mask = cv2.inRange(image, lower_white, upper_white)

        lower_black = np.array([0, 0, 0])
        upper_black = np.array([50, 50, 50])
        black_mask = cv2.inRange(image, lower_black, upper_black)

        # Find contours for windows and walls
        window_contours, _ = cv2.findContours(white_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        wall_contours, _ = cv2.findContours(black_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # Detect edges for door using Canny edge detection
        threshold1 = 50
        threshold2 = 150
        edges = cv2.Canny(image, threshold1, threshold2)
        kernel = np.ones((5, 5), np.uint8)
        edges = cv2.dilate(edges, kernel, iterations=1)
        edges = cv2.erode(edges, kernel, iterations=1)
        
        edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
        cv2.drawContours(edges_colored, contours, -1, (255, 0, 0), 10)
        edges_gray = cv2.cvtColor(edges_colored, cv2.COLOR_BGR2GRAY)
        
        contours, _ = cv2.findContours(edges_gray, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

        door_contours = []
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            perimeter = cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, 0.04 * perimeter, True)
            if len(approx) > 3 and len(approx) < 6 and 0.75 < w / float(h) < 1.25 and 100 < cv2.contourArea(contour) < 1500:
                door_contours.append(contour)

        # Morphological operations and re-detect contours
        white_mask = cv2.morphologyEx(white_mask, cv2.MORPH_OPEN, kernel)
        white_mask = cv2.morphologyEx(white_mask, cv2.MORPH_CLOSE, kernel)
        black_mask = cv2.morphologyEx(black_mask, cv2.MORPH_OPEN, kernel)
        black_mask = cv2.morphologyEx(black_mask, cv2.MORPH_CLOSE, kernel)
        window_contours, _ = cv2.findContours(white_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        wall_contours, _ = cv2.findContours(black_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # Create door-only image
        door_only_image = np.zeros_like(image)
        cv2.drawContours(door_only_image, door_contours, -1, (0, 255, 0), 1)
        
        
        # cv2.imshow('door_only_image', door_only_image)
        # cv2.imshow('edges_with_thickness_10', edges_colored)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

        #Skeletonize door area
        gray_wall_only2 = cv2.cvtColor(door_only_image, cv2.COLOR_BGR2GRAY)
        _, binary_image2 = cv2.threshold(gray_wall_only2, 127, 255, cv2.THRESH_BINARY)
        skeleton2 = skeletonize(binary_image2 // 255)
        skeleton2 = (skeleton2 * 255).astype(np.uint8)
        skeleton_image2 = np.zeros_like(image)
        skeleton_image2[skeleton2 == 255] = [255, 255, 255]
        skeleton_contours2, _ = cv2.findContours(skeleton2, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        linedoor = []
        for i, contour in enumerate(skeleton_contours2):
            for j in range(len(contour)):
                x1, y1 = contour[j][0]
                x2, y2 = contour[(j + 1) % len(contour)][0]
                linedoor.append((1, x1, y1, x2, y2))

        # Skeletonize wall area
        gray_wall_only = cv2.cvtColor(detected_contours, cv2.COLOR_BGR2GRAY)
        _, binary_image = cv2.threshold(gray_wall_only, 127, 255, cv2.THRESH_BINARY)
        skeleton = skeletonize(binary_image // 255)
        skeleton = (skeleton * 255).astype(np.uint8)
        skeleton_image = np.zeros_like(image)
        skeleton_image[skeleton == 255] = [255, 255, 255]
        skeleton_contours, _ = cv2.findContours(skeleton, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        linewall = []
        for i, contour in enumerate(skeleton_contours):
            for j in range(len(contour)):
                x1, y1 = contour[j][0]
                x2, y2 = contour[(j + 1) % len(contour)][0]
                linewall.append((0, x1, y1, x2, y2))
        result['wall_contours'] = linewall
        result['door_contours'] = linedoor
        return result

    """5.gradio接口返回这些需要的字段"""
    def process_urls(self,url):
        file_name = self.extract_file_name(url)
        wireframes_url = self.generate_urls(file_name)
        image = self.download_image(wireframes_url)
        processed_result = self.wireframes_image(image)
        processed_result2 = self.wireframes_image(image,2)
        wall_contours = processed_result2['wall_contours']
        door_contours = processed_result2['door_contours']
        wall_contours_1 = processed_result['wall_contours']
    
        return wall_contours, door_contours, wall_contours_1
        #return floorplan_rgb, cropped_image_rgb, first_row_sum, last_row_sum, max_min, wall_contours, door_contours

    """6.图片转换"""
    def image_to_base64(self,image):
        _, buffer = cv2.imencode('.jpg', image)
        img_base64 = base64.b64encode(buffer).decode('utf-8')
        return img_base64

# url = 'https://ke-image.ljcdn.com/hdic-frame/standard_477d6cbd-9d4f-42cf-b819-4082813fc63a.png.1440x1080.jpg?from=ke.com'
# url = 'http://ke-image.ljcdn.com/hdic-frame/standard_e520a0b7-4acb-4a0f-a7f6-caf64a8781cd.png.1440x1080.jpg?from=ke.com'
# url= "http://ke-image.ljcdn.com/hdic-frame/standard_18125c34-a482-48d1-bfd2-da1808ea979d.png.1440x1080.jpg?from=ke.com"
# url = "http://ke-image.ljcdn.com/hdic-frame/standard_472f2a3e-d695-4922-84de-d7aabac8e31d.png.1440x1080.jpg?from=ke.com"
# url = "http://ke-image.ljcdn.com/hdic-frame/standard_48b8d26f-3ab1-46a8-a8d3-cd34a3b1264f.png.1440x1080.jpg?from=ke.com"
# # url = "https://ke-image.ljcdn.com/hdic-frame/standard_172a6a46-b163-4794-884b-b98ff07129b4.png.1440x1080.jpg?from=ke.com"
# #多门
# # url = "https://ke-image.ljcdn.com/hdic-frame/standard_9e2aa60c-75fc-436c-bd87-6ad5f1c756bc.png.1440x1080.jpg?from=ke.com"
# url = "http://ke-image.ljcdn.com/hdic-frame/standard_1590e803-044e-4058-9728-0db7ebacff14.png.1440x1080.jpg?from=ke.com"
# url = "http://ke-image.ljcdn.com/hdic-frame/standard_45872acf-107e-494c-9820-a123a013c73a.png.1440x1080.jpg?from=ke.com"
# url = "https://ke-image.ljcdn.com/hdic-frame/standard_6adfce13-91bb-4bcc-9bd2-d363d856e965.png!m_fill,w_1000,h_750,l_bk,f_jpg,ls_50?from=ke.com"
# url = "https://ke-image.ljcdn.com/hdic-frame/standard_b9529263-6b9d-4af9-b4a2-1e52dabf56f6.png!m_fill,w_1000,h_750,l_bk,f_jpg,ls_50?from=ke.com"
# wireFrames_fun = WireFrames()
# wall_contours, door_contours, wall_contours_1 = wireFrames_fun.process_urls(url)
# print(f"wall_contours: {wall_contours}")
# print(f"door_contours: {door_contours}")